from tqdm import tqdm
import pandas as pd
import numpy as np
import torch
import cv2
import os
from PIL import Image
import warnings
from torchvision import transforms
from torchvision.models.feature_extraction import create_feature_extractor
import h5py
from torch.utils.data import Dataset




class Mydataset(Dataset):
    # 类初始化
    def __init__(self, img_paths, labels, transform_class, transform_defect):
        self.img_paths = img_paths
        self.labels = labels
        self.transform_class = transform_class
        self.transform_defect = transform_defect

    # 进行切片
    def __getitem__(self, index): #根据给出的索引进行切片，并对其进行数据处理转换成Tensor，返回成Tensor
        img_path = self.img_paths[index]
        label = self.labels[index]
        # get images from h5file
        f = h5py.File(img_path, 'r')  # 打开h5文件
        class_tiff = f['class_tiff'][:]
        defect_tiff = f['defect_tiff'][:]
        reference_tiff = f['reference_tiff'][:]
        f.close()

        class_tiff = Image.fromarray(class_tiff)
        defect_tiff = Image.fromarray(defect_tiff)
        reference_tiff = Image.fromarray(reference_tiff)

        if self.transform_class is not None:
            class_tiff = self.transform_class(class_tiff)

        if self.transform_defect is not None:
            defect_tiff = self.transform_defect(defect_tiff)
            reference_tiff = self.transform_defect(reference_tiff)
        sample = {'class_tiff': class_tiff, 'defect_tiff': defect_tiff, 'reference_tiff': reference_tiff, 'label': label}
        return sample
    # 返回长度
    def __len__(self):
        return len(self.img_paths)


def main():
    warnings.filterwarnings("ignore")
    device = torch.device('cuda:0' if torch.cuda.is_available() else 'cpu')
    weight = "model_123_92.304.pth"
    model = torch.load(weight)
    model = model.eval().to(device)

    def hook_fn(model, input, output):
        global hook_output
        hook_output = output

    encoding_array = []
    label_array = []

    categories = ['ADI_particle_developed', 'Array_peeling', 'Cu_missing', 'Other_peeling', 'Partial_etch',
                  'Pattern_fail',
                  'PR_peeling', 'Seam', 'Reference', 'Surface_particle', 'Burried_particle', 'Cu_diffuse',
                  'Prelayer_defect_developed',
                  'Void', 'Residue', 'Scratch']

    categories_to_id = dict((c, i) for i, c in enumerate(categories))
    id_to_categories = dict((v, k) for k, v in categories_to_id.items())

    def get_imgs_labels(dir):
        imgs = []
        labels = []
        categories = os.listdir(dir)
        for category in categories:
            category_path = os.path.join(dir, category)
            imgs_path = os.listdir(category_path)
            for img in imgs_path:
                imgs.append(os.path.join(category_path, img))
                labels.append(categories_to_id[category])
        return imgs, labels

    # 数据预处理
    transform_test = transforms.Compose([
        transforms.Resize((224, 224)),
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.40272543], std=[0.13901867])
    ])  # ([0.42368844, 0.42368844, 0.42368844], [0.14975642, 0.14975642, 0.14975642])


    resize_defect_ref_test = transforms.Compose([
        transforms.Resize((224, 224)),  # todo
        transforms.ToTensor(),
        transforms.Normalize(mean=[0.402817], std=[0.14428316])  # TODO
    ])

    val_imgs, val_labels = get_imgs_labels(r'/home/yeadc/Documents/cxj/datasets_h5_v2/val')
    dataset_test = Mydataset(val_imgs, val_labels, transform_test, resize_defect_ref_test)
    test_loader = torch.utils.data.DataLoader(dataset_test, batch_size=1, shuffle=False)

    for sample in test_loader:
        class_image = sample['class_tiff']
        defect_image = sample['defect_tiff']
        target = sample['label']
        class_image, defect_image = (class_image.to(device, non_blocking=True),
                                                                 defect_image.to(device, non_blocking=True))
        hook = model.pooling.register_forward_hook(hook_fn)
        output = model(class_image, defect_image)

        feature = hook_output.squeeze().detach().cpu().numpy()
        label = target.squeeze().detach().cpu().numpy()
        hook.remove()
        print(feature.shape)
        encoding_array.append(feature)
        label_array.append(label)
    np.save('feature_test.npy', encoding_array)
    np.save('label.npy', label_array)






if __name__ == '__main__':
    main()

